package com.flink.hbase.doris;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule;
import com.flink.hbase.doris.config.DorisConnectionConfig;
import com.flink.hbase.doris.model.UserEvent;
import org.apache.doris.flink.cfg.DorisExecutionOptions;
import org.apache.doris.flink.cfg.DorisOptions;
import org.apache.doris.flink.cfg.DorisReadOptions;
import org.apache.doris.flink.deserialization.DorisDeserializationSchema;
import org.apache.doris.flink.sink.DorisSink;
import org.apache.doris.flink.sink.writer.SimpleStringSerializer;
import org.apache.doris.flink.source.DorisSource;
import org.apache.doris.flink.source.DorisSourceBuilder;
import org.apache.flink.api.common.functions.FilterFunction;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.util.Collector;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.List;
import java.util.Properties;

/**
 * Flink-Doris StreamLoader 高效数据读写作业
 * 演示如何使用 Flink-Doris 连接器高效读取和写入 Doris 数据
 */
public class FlinkDorisStreamLoaderJob {
    
    private static final Logger LOG = LoggerFactory.getLogger(FlinkDorisStreamLoaderJob.class);
    private static final ObjectMapper objectMapper = new ObjectMapper();
    
    static {
        objectMapper.registerModule(new JavaTimeModule());
    }
    
    public static void main(String[] args) throws Exception {
        
        // 创建 Flink 执行环境
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(4);
        
        // 启用检查点
        env.enableCheckpointing(60000); // 每60秒创建一次检查点
        
        LOG.info("开始启动 Flink-Doris StreamLoader 作业");
        
        // 创建源表配置
        DorisConnectionConfig sourceConfig = createSourceConfig();
        
        // 从 Doris 读取数据流
        DataStream<UserEvent> sourceStream = createDorisSource(env, sourceConfig);
        
        // 数据处理和转换
        DataStream<UserEvent> processedStream = processUserEvents(sourceStream);
        
        // 创建目标表配置
        DorisConnectionConfig sinkConfig = createSinkConfig();
        
        // 写入到 Doris 目标表
        addDorisSink(processedStream, sinkConfig);
        
        // 执行作业
        env.execute("Flink-Doris StreamLoader 高效数据处理作业");
    }
    
    /**
     * 创建源表配置
     */
    private static DorisConnectionConfig createSourceConfig() {
        DorisConnectionConfig config = new DorisConnectionConfig(
            "127.0.0.1:8030",    // FE nodes
            "root",               // username
            "",                   // password
            "example_db",         // database
            "user_events"         // source table
        );
        
        // 优化读取参数
        config.setReadBatchSize(2048);      // 增大批次大小提高吞吐量
        config.setRequestRetries(5);        // 增加重试次数提高稳定性
        config.setRequestConnectTimeoutMs(60000);  // 连接超时
        config.setRequestReadTimeoutMs(60000);     // 读取超时
        config.setRequestQueryTimeoutS(7200);      // 查询超时
        
        LOG.info("源表配置: {}", config);
        return config;
    }
    
    /**
     * 创建目标表配置
     */
    private static DorisConnectionConfig createSinkConfig() {
        DorisConnectionConfig config = new DorisConnectionConfig(
            "127.0.0.1:8030",       // FE nodes
            "root",                 // username
            "",                     // password
            "example_db",           // database
            "processed_events"      // target table
        );
        
        // 优化写入参数
        config.setBatchSize(2000);          // 批次大小
        config.setBatchIntervalMs(5000L);   // 批次间隔
        config.setMaxRetries(3);            // 最大重试次数
        config.setStreamLoadFormat("json"); // 使用JSON格式
        
        LOG.info("目标表配置: {}", config);
        return config;
    }
    
    /**
     * 创建 Doris 数据源
     */
    public static DataStream<UserEvent> createDorisSource(
            StreamExecutionEnvironment env,
            DorisConnectionConfig config) {
        
        LOG.info("创建 Doris Source，连接到: {}", config.getFenodes());
        
        // Doris 连接选项
        DorisOptions.Builder dorisBuilder = DorisOptions.builder()
                .setFenodes(config.getFenodes())
                .setTableIdentifier(config.getDatabase() + "." + config.getTable())
                .setUsername(config.getUsername())
                .setPassword(config.getPassword());
        
        DorisOptions dorisOptions = dorisBuilder.build();
        
        // Doris 读取选项 - 优化配置
        DorisReadOptions readOptions = DorisReadOptions.builder()
                .setDeserializeArrowAsync(false)      // 禁用异步Arrow反序列化
                .setDeserializeQueueSize(64)          // 反序列化队列大小
                .setExecMemLimit(2147483648L)         // 执行内存限制 2GB
                .setRequestTabletSize(Integer.MAX_VALUE) // Tablet大小
                .setRequestConnectTimeoutMs(config.getRequestConnectTimeoutMs())
                .setRequestReadTimeoutMs(config.getRequestReadTimeoutMs())
                .setRequestQueryTimeoutS(config.getRequestQueryTimeoutS())
                .setRequestRetries(config.getRequestRetries())
                .setRequestBatchSize(config.getReadBatchSize())
                .build();
        
        // 创建 Doris Source
        DorisSource<UserEvent> dorisSource = DorisSourceBuilder.<UserEvent>builder()
                .setDorisReadOptions(readOptions)
                .setDorisOptions(dorisOptions)
                .setDeserializer(new UserEventDeserializationSchema())
                .build();
        
        return env.fromSource(dorisSource,
                org.apache.flink.api.common.eventtime.WatermarkStrategy.noWatermarks(),
                "doris-source")
                .returns(TypeInformation.of(UserEvent.class));
    }
    
    /**
     * 处理用户事件数据流
     */
    private static DataStream<UserEvent> processUserEvents(DataStream<UserEvent> sourceStream) {
        LOG.info("开始处理用户事件数据流");
        
        return sourceStream
                // 过滤无效数据
                .filter(new FilterFunction<UserEvent>() {
                    @Override
                    public boolean filter(UserEvent event) throws Exception {
                        return event != null 
                               && event.getUserId() != null 
                               && event.getUserId() > 0
                               && event.getEventTime() != null;
                    }
                })
                // 数据清洗和标准化
                .map(new MapFunction<UserEvent, UserEvent>() {
                    @Override
                    public UserEvent map(UserEvent event) throws Exception {
                        // 设备类型标准化
                        if (event.getDeviceType() == null || event.getDeviceType().trim().isEmpty()) {
                            event.setDeviceType("unknown");
                        }
                        
                        // PV值标准化
                        if (event.getPv() == null || event.getPv() <= 0) {
                            event.setPv(1L);
                        }
                        
                        // 时长标准化
                        if (event.getDuration() == null || event.getDuration() < 0) {
                            event.getDuration();
                        }
                        
                        // IP地址脱敏处理
                        if (event.getIpAddress() != null) {
                            String[] parts = event.getIpAddress().split("\\.");
                            if (parts.length == 4) {
                                event.setIpAddress(parts[0] + "." + parts[1] + ".*.* ");
                            }
                        }
                        
                        return event;
                    }
                })
                .name("process-user-events");
    }
    
    /**
     * 添加 Doris Sink
     */
    public static void addDorisSink(DataStream<UserEvent> dataStream, DorisConnectionConfig config) {
        LOG.info("添加 Doris Sink，写入到: {}.{}", config.getDatabase(), config.getTable());
        
        // Doris 连接选项
        DorisOptions dorisOptions = DorisOptions.builder()
                .setFenodes(config.getFenodes())
                .setTableIdentifier(config.getDatabase() + "." + config.getTable())
                .setUsername(config.getUsername())
                .setPassword(config.getPassword())
                .build();
        
        // StreamLoad 属性配置
        Properties streamLoadProps = new Properties();
        streamLoadProps.setProperty("format", config.getStreamLoadFormat());
        streamLoadProps.setProperty("read_json_by_line", "true");
        streamLoadProps.setProperty("load_to_single_tablet", "false");
        streamLoadProps.setProperty("timeout", "600");
        streamLoadProps.setProperty("max_filter_ratio", "0.1");
        
        // Doris 执行选项
        DorisExecutionOptions executionOptions = DorisExecutionOptions.builder()
                .setLabelPrefix("flink-doris-streamloader")
                .setStreamLoadProp(streamLoadProps)
                .setBatchSize(config.getBatchSize())
                .setBatchIntervalMs(config.getBatchIntervalMs())
                .setMaxRetries(config.getMaxRetries())
                .build();
        
        // 创建 Doris Sink
        DorisSink.Builder<UserEvent> sinkBuilder = DorisSink.<UserEvent>builder()
                .setDorisReadOptions(DorisReadOptions.builder().build())
                .setDorisOptions(dorisOptions)
                .setDorisExecutionOptions(executionOptions)
                .setSerializer(new UserEventStringSerializer());
        
        dataStream.sinkTo(sinkBuilder.build()).name("doris-sink");
    }
    
    /**
     * 用户事件反序列化 Schema
     */
    public static class UserEventDeserializationSchema implements DorisDeserializationSchema<UserEvent> {
        
        private static final DateTimeFormatter dateTimeFormatter = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss");
        
        @Override
        public void deserialize(List<String> record, Collector<UserEvent> out) throws IOException {
            if (record == null || record.size() < 12) {
                LOG.warn("记录字段不足，跳过: {}", record);
                return;
            }
            
            try {
                UserEvent userEvent = new UserEvent();
                
                // 解析各个字段
                userEvent.setEventTime(LocalDateTime.parse(record.get(0), dateTimeFormatter));
                userEvent.setUserId(Long.parseLong(record.get(1)));
                userEvent.setEventType(record.get(2));
                userEvent.setPageUrl(record.get(3));
                userEvent.setDeviceType(record.get(4));
                userEvent.setBrowser(record.get(5));
                userEvent.setOs(record.get(6));
                userEvent.setIpAddress(record.get(7));
                userEvent.setCountry(record.get(8));
                userEvent.setCity(record.get(9));
                userEvent.setPv(Long.parseLong(record.get(10)));
                userEvent.setDuration(Long.parseLong(record.get(11)));
                
                out.collect(userEvent);
                
            } catch (Exception e) {
                LOG.error("解析记录失败: {}, 错误: {}", record, e.getMessage());
                // 继续处理下一条记录
            }
        }
        
        @Override
        public TypeInformation<UserEvent> getProducedType() {
            return TypeInformation.of(UserEvent.class);
        }
    }
    
    /**
     * 用户事件字符串序列化器
     */
    private static class UserEventStringSerializer extends SimpleStringSerializer {
        
        @Override
        public String serialize(Object element) {
            try {
                if (element instanceof UserEvent) {
                    return objectMapper.writeValueAsString(element);
                }
                return element.toString();
            } catch (Exception e) {
                LOG.error("序列化 UserEvent 失败", e);
                throw new RuntimeException("序列化失败", e);
            }
        }
    }
} 